import os, sys
from socket import *
import time

class FtpClient(object):
    def __init__(self, sockfd):
        self.sockfd = sockfd
    
    def do_list(self):
        self.sockfd.send(b'L')
        data = self.sockfd.recv(128).decode()
        if data == 'OK':
            while True:
                file = self.sockfd.recv(1024).decode()
                if file == '##':
                    break
                print(file)
        else:
            print(data)
    
    def do_quit(self):
        self.sockfd.send(b'Q')
        self.sockfd.close()
        sys.exit('Thank you for using')

    def do_get(self, filename):
        self.sockfd.send(('G ' + filename).encode())
        data = self.sockfd.recv(128).decode()
        if data == 'OK':
            fd = open(filename, 'wb')
            while True:
                data = self.sockfd.recv(1024)
                if data == b'##':
                    break
                fd.write(data)
            fd.close()
            print('%sdownload success!' % filename)
        else:
            print(data)
    
    def do_put(self, filename):
        try:
            f = open(filename, 'rb')
        except:
            print('File is not exist')
            return

        self.sockfd.send(('P ' + filename).encode())
        data = self.sockfd.recv(128).decode()
        if data == 'OK':
            while True:
                data = f.read(1024)
                if not data:
                    time.sleep(0.1)
                    self.sockfd.send(b'##')
                self.sockfd.send(data)
                break
            f.close()
            print('Transfer success!')
        else:
            print(data)


def main():
    if len(sys.argv) < 3:
        print('argv is error')
        return
    HOST = sys.argv[1]
    PORT = int(sys.argv[2])
    ADDR = (HOST, PORT)

    sockfd = socket()
    try:
        sockfd.connect(ADDR)
    except:
        print('Server-connection failed!')
        return
    
    ftp = FtpClient(sockfd)
    while True:
        print('\n======commend selections======')
        print('***         list           ***')
        print('***       get file         ***')
        print('***       put file         ***')
        print('***         quit           ***')
        print('==============================\n')

        cmd = input('input your commend >>')
        if cmd.strip() == 'list':
            ftp.do_list()
        elif cmd.strip() == 'quit':
            ftp.do_quit()
        elif cmd[:3] == 'get':
            filename = cmd.split(' ')[-1]
            ftp.do_get(filename)
        elif cmd[:3] == 'put':
            filename = cmd.split(' ')[-1]
            ftp.do_put(filename)

        else:
            print('the wrong commend!')

main()  
